package com.hanliukui.ai.langchain4j.infrastructure;

import com.hanliukui.ai.langchain4j.infrastructure.domain.po.ChatMessagesPo;
import com.hanliukui.ai.langchain4j.infrastructure.domain.repository.ChatMessagesPoRepository;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.ChatMessageDeserializer;
import dev.langchain4j.data.message.ChatMessageSerializer;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;

import java.util.List;

public class JpaChatMemoryStore implements ChatMemoryStore {

    private final ChatMessagesPoRepository repository;

    public JpaChatMemoryStore(ChatMessagesPoRepository repository) {
        this.repository = repository;
    }

    @Override
    public List<ChatMessage> getMessages(Object memoryId) {
        ChatMessagesPo chatMessagesPo = repository.findByMemoryId((String) memoryId);
        if (chatMessagesPo != null) {
            return ChatMessageDeserializer.messagesFromJson(chatMessagesPo.getContent());
        }
        return List.of();
    }

    @Override
    public void updateMessages(Object memoryId, List<ChatMessage> list) {
        ChatMessagesPo chatMessagesPo = repository.findByMemoryId((String) memoryId);
        if (chatMessagesPo == null) {
            chatMessagesPo = new ChatMessagesPo();
            chatMessagesPo.setMemoryId((String) memoryId);
        }
        chatMessagesPo.setContent(ChatMessageSerializer.messagesToJson(list));
        repository.save(chatMessagesPo);
    }

    @Override
    public void deleteMessages(Object memoryId) {
        repository.deleteByMemoryId((String) memoryId);
    }
}
